import threading
import time
import psutil
from datetime import datetime
from collections import deque
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from flask import Flask, render_template, jsonify, request
import io
import base64
import json
import os

app = Flask(__name__)

# 全局变量存储流量数据
traffic_data = {
    'incoming': deque(maxlen=300),  # 存储最近5分钟的入站流量(每秒一个数据点)
    'outgoing': deque(maxlen=300),  # 存储最近5分钟的出站流量
    'timestamps': deque(maxlen=300)  # 存储对应的时间戳
}


# 流量统计类
class TrafficMonitor(threading.Thread):
    def __init__(self, interface_name=None):
        super().__init__()
        self.daemon = True
        self.last_incoming = 0
        self.last_outgoing = 0
        self.interface_name = interface_name
        self.running = True

    def run(self):
        print("开始监控网络流量...")

        # 获取网络接口
        interfaces = psutil.net_io_counters(pernic=True)

        # 如果没有指定接口，使用第一个活动接口
        if not self.interface_name:
            for name in interfaces:
                if name != 'lo' and interfaces[name].bytes_sent > 0:
                    self.interface_name = name
                    break

        if not self.interface_name:
            print("未找到可用的网络接口")
            return

        print(f"监控接口: {self.interface_name}")

        # 初始化计数器
        self.last_incoming = interfaces[self.interface_name].bytes_recv
        self.last_outgoing = interfaces[self.interface_name].bytes_sent

        # 开始监控循环
        while self.running:
            try:
                # 获取当前流量统计
                current_stats = psutil.net_io_counters(pernic=True)[self.interface_name]
                current_incoming = current_stats.bytes_recv
                current_outgoing = current_stats.bytes_sent

                # 计算每秒流量
                incoming_rate = current_incoming - self.last_incoming
                outgoing_rate = current_outgoing - self.last_outgoing

                # 更新计数器
                self.last_incoming = current_incoming
                self.last_outgoing = current_outgoing

                # 更新全局流量数据
                now = datetime.now()
                traffic_data['timestamps'].append(now)
                traffic_data['incoming'].append(incoming_rate)
                traffic_data['outgoing'].append(outgoing_rate)

                # 每秒更新一次
                time.sleep(1)

            except Exception as e:
                print(f"监控出错: {e}")
                time.sleep(5)

    def stop(self):
        self.running = False


# 创建并启动流量监控线程
monitor = TrafficMonitor()


# Flask路由
@app.route('/')
def index():
    return render_template('index.html')


@app.route('/traffic-data')
def get_traffic_data():
    # 返回JSON格式的流量数据
    data = {
        'timestamps': [ts.strftime('%H:%M:%S') for ts in traffic_data['timestamps']],
        'incoming': list(traffic_data['incoming']),
        'outgoing': list(traffic_data['outgoing'])
    }
    return jsonify(data)


@app.route('/total-traffic')
def get_total_traffic():
    # 获取总流量统计
    interfaces = psutil.net_io_counters(pernic=True)
    interface_name = monitor.interface_name

    total_incoming = interfaces[interface_name].bytes_recv if interface_name in interfaces else 0
    total_outgoing = interfaces[interface_name].bytes_sent if interface_name in interfaces else 0

    return jsonify({
        'total_incoming': total_incoming,
        'total_outgoing': total_outgoing
    })


@app.route('/traffic-plot')
def get_traffic_plot():
    # 生成流量图表并返回base64编码的图像
    if not traffic_data['timestamps']:
        return "暂无数据"

    plt.figure(figsize=(10, 6))

    # 转换数据为KB/s
    incoming_kb = [x / 1024 for x in traffic_data['incoming']]
    outgoing_kb = [x / 1024 for x in traffic_data['outgoing']]

    plt.plot(traffic_data['timestamps'], incoming_kb, label='入站流量 (KB/s)')
    plt.plot(traffic_data['timestamps'], outgoing_kb, label='出站流量 (KB/s)')

    # 格式化图表
    plt.xlabel('时间')
    plt.ylabel('流量 (KB/秒)')
    plt.title('实时网络流量监控')
    plt.legend()
    plt.grid(True)
    plt.gcf().autofmt_xdate()
    plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%H:%M:%S'))

    # 将图表转换为base64编码
    img = io.BytesIO()
    plt.savefig(img, format='png')
    img.seek(0)
    plot_url = base64.b64encode(img.getvalue()).decode()
    plt.close()

    return f'<img src="data:image/png;base64,{plot_url}">'


@app.route('/report')
def generate_report():
    # 生成带宽使用报告
    if not traffic_data['timestamps']:
        return "<p>暂无数据可生成报告</p>"

    # 计算统计信息
    total_in = sum(traffic_data['incoming'])
    total_out = sum(traffic_data['outgoing'])
    avg_in = total_in / len(traffic_data['incoming'])
    avg_out = total_out / len(traffic_data['outgoing'])
    max_in = max(traffic_data['incoming'])
    max_out = max(traffic_data['outgoing'])

    # 转换为更友好的单位
    def format_bytes(bytes):
        for unit in ['B', 'KB', 'MB', 'GB']:
            if bytes < 1024.0:
                return f"{bytes:.2f} {unit}"
            bytes /= 1024.0
        return f"{bytes:.2f} TB"

    def format_bps(bytes_per_sec):
        return format_bytes(bytes_per_sec) + "/s"

    report = f"""
    <h2>带宽使用报告</h2>
    <p>统计时间段: {traffic_data['timestamps'][0]} 到 {traffic_data['timestamps'][-1]}</p>
    <p>监控时长: {len(traffic_data['timestamps'])} 秒</p>
    <p>总入站流量: {format_bytes(total_in)}</p>
    <p>总出站流量: {format_bytes(total_out)}</p>
    <p>平均入站速率: {format_bps(avg_in)}</p>
    <p>平均出站速率: {format_bps(avg_out)}</p>
    <p>最大入站速率: {format_bps(max_in)}</p>
    <p>最大出站速率: {format_bps(max_out)}</p>
    """

    return report


@app.route('/download-report')
def download_report():
    # 生成并下载JSON格式的详细报告
    if not traffic_data['timestamps']:
        return "暂无数据", 404

    # 准备报告数据
    report_data = {
        "generated_at": datetime.now().isoformat(),
        "time_period": {
            "start": traffic_data['timestamps'][0].isoformat() if traffic_data['timestamps'] else None,
            "end": traffic_data['timestamps'][-1].isoformat() if traffic_data['timestamps'] else None,
            "duration_seconds": len(traffic_data['timestamps'])
        },
        "traffic_data": {
            "timestamps": [ts.isoformat() for ts in traffic_data['timestamps']],
            "incoming_bytes_per_sec": list(traffic_data['incoming']),
            "outgoing_bytes_per_sec": list(traffic_data['outgoing'])
        },
        "statistics": {
            "total_incoming_bytes": sum(traffic_data['incoming']),
            "total_outgoing_bytes": sum(traffic_data['outgoing']),
            "avg_incoming_bytes_per_sec": sum(traffic_data['incoming']) / len(traffic_data['incoming']),
            "avg_outgoing_bytes_per_sec": sum(traffic_data['outgoing']) / len(traffic_data['outgoing']),
            "max_incoming_bytes_per_sec": max(traffic_data['incoming']) if traffic_data['incoming'] else 0,
            "max_outgoing_bytes_per_sec": max(traffic_data['outgoing']) if traffic_data['outgoing'] else 0
        }
    }

    # 转换为JSON字符串
    report_json = json.dumps(report_data, indent=2)

    # 创建响应
    from flask import Response
    response = Response(
        report_json,
        mimetype="application/json",
        headers={"Content-Disposition": "attachment;filename=bandwidth_report.json"}
    )

    return response


@app.route('/change-view')
def change_view():
    # 改变数据视图范围
    range = request.args.get('range', 'minute')

    # 根据范围调整数据保留数量
    if range == 'minute':
        new_maxlen = 60  # 1分钟
    elif range == 'hour':
        new_maxlen = 3600  # 1小时
    elif range == 'day':
        new_maxlen = 86400  # 24小时
    else:
        new_maxlen = 300  # 默认5分钟

    # 创建新的deque并复制现有数据
    def resize_deque(old_deque, new_maxlen):
        new_deque = deque(maxlen=new_maxlen)
        for item in old_deque:
            new_deque.append(item)
        return new_deque

    traffic_data['incoming'] = resize_deque(traffic_data['incoming'], new_maxlen)
    traffic_data['outgoing'] = resize_deque(traffic_data['outgoing'], new_maxlen)
    traffic_data['timestamps'] = resize_deque(traffic_data['timestamps'], new_maxlen)

    return jsonify({"status": "success", "new_maxlen": new_maxlen})


if __name__ == '__main__':
    # 启动流量监控线程
    monitor.start()

    # 启动Flask应用
    print("启动带宽监控Web界面...")
    print("请访问 http://localhost:5000")
    app.run(debug=True, host='0.0.0.0', port=5000, use_reloader=False)
